"""
SSH远程控制工具中继服务器API路由
"""

from flask import Blueprint, request, jsonify
import sqlite3
import random
import string
import datetime
import os

api_bp = Blueprint("api", __name__, url_prefix="/api")

# 数据库路径
DB_PATH = os.path.join(os.path.dirname(os.path.dirname(__file__)), "database", "app.db")

def init_db():
    """初始化数据库"""
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    
    # 创建连接表
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS connections (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            key TEXT UNIQUE NOT NULL,
            client_info TEXT,
            ssh_port INTEGER DEFAULT 22,
            tunnel_port INTEGER,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            expires_at TIMESTAMP,
            last_heartbeat TIMESTAMP,
            status TEXT DEFAULT "active"
        )
    """)
    
    conn.commit()
    conn.close()

def generate_key():
    """生成8位随机密钥"""
    return "".join(random.choices(string.ascii_uppercase + string.digits, k=8))

def generate_tunnel_port():
    """生成隧道端口"""
    return random.randint(8000, 9000)

@api_bp.route("/register", methods=["POST"])
def register():
    """注册被控端连接"""
    try:
        data = request.get_json() or {}
        client_info = data.get("client_info", "")
        ssh_port = data.get("ssh_port", 22)
        
        # 生成密钥和隧道端口
        key = generate_key()
        tunnel_port = generate_tunnel_port()
        
        # 设置过期时间（24小时后）
        expires_at = datetime.datetime.now() + datetime.timedelta(hours=24)
        
        # 保存到数据库
        conn = sqlite3.connect(DB_PATH)
        cursor = conn.cursor()
        
        cursor.execute("""
            INSERT INTO connections (key, client_info, ssh_port, tunnel_port, expires_at, last_heartbeat)
            VALUES (?, ?, ?, ?, ?, ?)
        """, (key, client_info, ssh_port, tunnel_port, expires_at, datetime.datetime.now()))
        
        conn.commit()
        conn.close()
        
        return jsonify({
            "status": "success",
            "key": key,
            "tunnel_port": tunnel_port,
            "expires_at": expires_at.isoformat()
        })
        
    except Exception as e:
        return jsonify({"status": "error", "message": str(e)}), 500

@api_bp.route("/connect", methods=["POST"])
def connect():
    """验证密钥并获取连接信息"""
    try:
        data = request.get_json() or {}
        key = data.get("key", "").upper()
        
        if not key:
            return jsonify({"status": "error", "message": "密钥不能为空"}), 400
        
        conn = sqlite3.connect(DB_PATH)
        cursor = conn.cursor()
        
        cursor.execute("""
            SELECT key, tunnel_port, ssh_port, expires_at, status
            FROM connections 
            WHERE key = ? AND status = "active" AND expires_at > datetime("now")
        """, (key,))
        
        result = cursor.fetchone()
        conn.close()
        
        if not result:
            return jsonify({"status": "error", "message": "密钥无效或已过期"}), 404
        
        return jsonify({
            "status": "success",
            "key": result[0],
            "tunnel_port": result[1],
            "ssh_port": result[2],
            "expires_at": result[3]
        })
        
    except Exception as e:
        return jsonify({"status": "error", "message": str(e)}), 500

@api_bp.route("/heartbeat", methods=["POST"])
def heartbeat():
    """发送心跳保持连接"""
    try:
        data = request.get_json() or {}
        key = data.get("key", "").upper()
        
        if not key:
            return jsonify({"status": "error", "message": "密钥不能为空"}), 400
        
        conn = sqlite3.connect(DB_PATH)
        cursor = conn.cursor()
        
        cursor.execute("""
            UPDATE connections 
            SET last_heartbeat = datetime("now")
            WHERE key = ? AND status = "active" AND expires_at > datetime("now")
        """, (key,))
        
        if cursor.rowcount == 0:
            conn.close()
            return jsonify({"status": "error", "message": "密钥无效或已过期"}), 404
        
        conn.commit()
        conn.close()
        
        return jsonify({"status": "success", "message": "心跳更新成功"})
        
    except Exception as e:
        return jsonify({"status": "error", "message": str(e)}), 500

@api_bp.route("/unregister", methods=["POST"])
def unregister():
    """注销连接"""
    try:
        data = request.get_json() or {}
        key = data.get("key", "").upper()
        
        if not key:
            return jsonify({"status": "error", "message": "密钥不能为空"}), 400
        
        conn = sqlite3.connect(DB_PATH)
        cursor = conn.cursor()
        
        cursor.execute("""
            UPDATE connections 
            SET status = "inactive"
            WHERE key = ?
        """, (key,))
        
        conn.commit()
        conn.close()
        
        return jsonify({"status": "success", "message": "连接已注销"})
        
    except Exception as e:
        return jsonify({"status": "error", "message": str(e)}), 500

@api_bp.route("/status", methods=["GET"])
def status():
    """获取服务器状态"""
    try:
        conn = sqlite3.connect(DB_PATH)
        cursor = conn.cursor()
        
        # 获取活跃连接数
        cursor.execute("""
            SELECT COUNT(*) FROM connections 
            WHERE status = "active" AND expires_at > datetime("now")
        """)
        active_connections = cursor.fetchone()[0]
        
        # 获取总连接数
        cursor.execute("SELECT COUNT(*) FROM connections")
        total_connections = cursor.fetchone()[0]
        
        conn.close()
        
        return jsonify({
            "status": "running",
            "active_connections": active_connections,
            "total_connections": total_connections,
            "server_time": datetime.datetime.now().isoformat()
        })
        
    except Exception as e:
        return jsonify({"status": "error", "message": str(e)}), 500

# 初始化数据库
init_db()



